import collections


def get_weighted_model(model_list, model_weights=None):
    '''
       模型权重按照weights加权得到
       model_list和model_weights中的权重一一对应
    '''
    if weights == None:
        weights = [1./len(model_list)] * len(model_list)
    worker_state_dict = [x.state_dict() for x in model_list]
    weight_keys = list(worker_state_dict[0].keys())
    fed_state_dict = collections.OrderedDict()
    for key in weight_keys:
        key_sum = 0
        for idx in range(len(model_list)):
            key_sum += worker_state_dict[idx][key] * model_weights[idx]
        fed_state_dict[key] = key_sum

    model_list[0].load_state_dict(fed_state_dict)
    return model_list[0]